#include <hal/ahci/ahci.h>
#include <hal/ahci/hba.h>
#include <hal/ahci/sata.h>

#include <ds/parted_buffer.h>

#include <arch/io.h>

#include <mm/valloc.h>
#include <mm/mm.h>

#include <string.h>
#include <common.h>

#define LOG_TAG "AHCI"
#include <log.h>

void hba_reset_port(uptr* port_reg) {
    port_reg[HBA_REG_PxCMD] &= ~HBA_PxCMD_ST;
    port_reg[HBA_REG_PxCMD] &= ~HBA_PxCMD_FRE;
    int cnt = wait_until_expire(!(port_reg[HBA_REG_PxCMD] & HBA_PxCMD_CR), 500000);
    if (cnt) {
        return;
    }
    // 如果port未响应，则继续执行重置
    port_reg[HBA_REG_PxSCTL] = (port_reg[HBA_REG_PxSCTL] & ~0xF) | 1;
    io_delay(100000);
    port_reg[HBA_REG_PxSCTL] &= ~0xF;
}

void hba_bind_buf(hba_cmdh_t* cmdh, hba_cmdt_t* cmdt, void* buf, size_t size) {
    assert_msg(size <= 0x400000, "HBA: Buffer too big");
    cmdh->prdt_len = 1;
    cmdt->entries[0] = (hba_prdte_t){
        .data_base = get_pa((uptr)buf),
        .byte_count = size - 1
    };
}

void hba_bind_parted_buf(hba_cmdh_t* cmdh, hba_cmdt_t* cmdt, pbuffer_t* pbuf) {
    size_t i = 0;
    pbuffer_t* pos = pbuf;

    do {
        assert_msg(i < HBA_MAX_PRDTE, "HBA: Too many PRDTEs");
        assert_msg(pos->size <= 0x400000, "HBA: Buffer too big");

        cmdt->entries[i++] = (hba_prdte_t){
            .data_base = (uptr)pos->paddr,
            .byte_count = pos->size - 1
        };
        pos = list_entry(pos->part.next, pbuffer_t, part);
    } while (pos != pbuf);

    cmdh->prdt_len = i + 1;
}

int hba_init_device(hba_port_t* port) {
    port->regs[HBA_REG_PxIE] &= ~HBA_PxINTR_DHR;
    u16* data_in = valloc_dma(512);

    hba_cmdh_t* cmdh;
    hba_cmdt_t* cmdt;
    int slot = hba_prepare_cmd(port, &cmdt, &cmdh);
    hba_bind_buf(cmdh, cmdt, data_in, 512);

    port->device = valloc(sizeof(hba_dev_t));
    memset(port->device, 0, sizeof(hba_dev_t));
    port->device->port = port;

    sata_cmd_fis_t* cmd_fis = (sata_cmd_fis_t*)cmdt->command_fis;

    if (port->regs[HBA_REG_PxSIG] == HBA_DEV_SIG_ATA) {
        // hard disk
        sata_create_fis(cmd_fis, ATA_IDENTIFY_DEVICE, 0, 0);
    } else {
        // CD-ROM or other device
        LOGW("We dont support CD-ROM now!\n");
        goto fail;
    }

    if (!ahci_try_send(port, slot)) {
        LOGE("Failed to send cmd\n");
        goto fail;
    }

    ahci_parse_dev_info(port->device, data_in);

    if ((port->device->flags & HBA_DEV_FATAPI)) {
        LOGE("ATAPI device is not supported yet\n");
        goto fail;
    }

    port->regs[HBA_REG_PxIE] |= HBA_PxINTR_DHR | HBA_PxINTR_OF | HBA_PxINTR_TFE;
    port->device->ops.submit = sata_submit;

    vfree_dma(data_in);
    vfree_dma(cmdt);
    return 1;

fail:
    port->regs[HBA_REG_PxIE] |= HBA_PxINTR_DHR | HBA_PxINTR_OF | HBA_PxINTR_TFE;
    vfree_dma(data_in);
    vfree_dma(cmdt);
    return 0;
}